'''
Author: Diana Tang
Date: 2025-04-05 01:40:24
LastEditors: Diana Tang
Description: some description
FilePath: /pub-scholar-assist/scripts/rebuild_vector_store.py
'''
#!/usr/bin/env python
"""
重建向量存储脚本
用于清除旧索引并重新构建知识库
"""

import os
import sys
import shutil
import argparse

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.model_loader import model_loader
from app.document_processor import DocumentProcessor
from app.vector_store import VectorStore

def rebuild_vector_store(docs_dir, config_path="config/config.yaml", force=False):
    """清除现有索引并重建向量存储"""
    
    # 确认操作
    if not force:
        confirm = input("这将删除现有的向量索引并重新构建。确定要继续吗? [y/N] ")
        if confirm.lower() not in ["y", "yes"]:
            print("操作已取消")
            return
    
    print(f"开始重建向量存储，使用文档目录: {docs_dir}")
    
    # 加载嵌入模型
    print("加载嵌入模型...")
    embedding_model = model_loader.load_embedding_model()
    
    # 初始化文档处理器
    doc_processor = DocumentProcessor(config_path)
    
    # 初始化向量存储
    vector_store = VectorStore(embedding_model, config_path)
    
    # 获取索引路径
    index_path = vector_store.index_path
    metadata_path = vector_store.metadata_path
    
    # 删除现有索引文件
    if os.path.exists(index_path):
        print(f"删除现有索引文件: {index_path}")
        os.remove(index_path)
    
    if os.path.exists(metadata_path):
        print(f"删除现有元数据文件: {metadata_path}")
        os.remove(metadata_path)
    
    # 确保目录存在
    os.makedirs(os.path.dirname(index_path), exist_ok=True)
    
    # 处理文档
    print(f"处理文档目录: {docs_dir}")
    chunks = doc_processor.process_directory(docs_dir)
    
    # 添加到向量存储
    print(f"将 {len(chunks)} 个文档块添加到向量存储...")
    vector_store.add_documents(chunks)
    
    # 保存索引
    print("保存索引...")
    vector_store.save()
    
    print(f"向量存储重建完成，共索引了 {len(chunks)} 个文档块")
    print(f"索引保存在: {index_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="重建向量存储")
    parser.add_argument("--docs", required=True, help="文档目录路径")
    parser.add_argument("--config", default="config/config.yaml", help="配置文件路径")
    parser.add_argument("--force", action="store_true", help="强制重建，不询问确认")
    
    args = parser.parse_args()
    
    rebuild_vector_store(args.docs, args.config, args.force)